import re
import json
import requests
import os
import shutil
import logging
import zipfile
from openai import OpenAI
import time

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()],
)

# 检查是否有config.json，没有则提醒用户创建。有则读取config.json
# 定义配置文件路径
config_file_path = "config.json"
# 检查文件是否存在
if not os.path.exists(config_file_path):
    logging.info("警告：未找到 config.json 文件，请先创建该文件！")
    # 提示用户如何创建文件（可选）
    logging.info("请参考README.MD创建 config.json 文件：")
else:
    try:
        # 尝试打开并读取 config.json 文件
        with open(config_file_path, "r", encoding="utf-8") as config_file:
            config = json.load(config_file)
            logging.info("配置文件已成功加载：")
            logging.info(config)
    except json.JSONDecodeError:
        # 如果文件内容不是有效的 JSON 格式
        logging.info("错误：config.json 文件格式不正确，请检查文件内容是否符合 JSON 格式！")
    except Exception as e:
        # 捕获其他可能的异常
        logging.info(f"发生未知错误：{e}")

# MinerU配置项
minerU_api_key = config.get("minerU_api_key", "")
header = {
    'Content-Type': 'application/json',
    "Authorization": f"Bearer {minerU_api_key}"
}
upload_url = 'https://mineru.net/api/v4/file-urls/batch'
query_url_template = 'https://mineru.net/api/v4/extract-results/batch/{}'
download_folder = 'results'
os.makedirs(download_folder, exist_ok=True)

# OpenAI 客户端配置
client = OpenAI(
    api_key= config.get('bailian_api_key', ""),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

def upload_pdf(file_name):
    """上传 PDF 文件到转换任务"""
    try:
        data = {'files': [{'name' : file_name, 'is_ocr' : True}], 'language' : 'en',"enable_formula": True,"enable_table": True}
        res = requests.post(upload_url, headers=header, json=data)
        task_info = res.json()
        if res.status_code != 200 or task_info.get('code') != 0:
            logging.error("Failed to create task")
            return None
        logging.info('response success. result:{}'.format(task_info))
        batch_id = task_info["data"]["batch_id"]
        urls = task_info["data"]["file_urls"]
        logging.info('batch_id:{},urls:{}'.format(batch_id, urls))
        with open(file_name, 'rb') as f:
            res_upload = requests.put(urls[0], data=f)
        if res_upload.status_code == 200:
            logging.info(f"Task created successfully with ID: {batch_id}")
        else:
            logging.info("Task created failed")
        return batch_id
    except Exception as e:
        logging.error(f"Error creating task: {e}")
        return None

def check_task_status(task_id):
    """轮询检查任务状态，直到完成或失败"""
    try:
        while True:
            query_url = query_url_template.format(task_id)
            res = requests.get(query_url, headers=header)
            status_info = res.json()
            state = status_info['data']['extract_result'][0].get('state')
            if state == 'done':
                logging.info("Task completed successfully.")
                return status_info['data']['extract_result'][0]['full_zip_url']
            elif state == 'failed':
                logging.error(f"Task failed: {status_info['data']['extract_result'][0].get('err_msg')}")
                return None
            else:
                logging.info(f"Current state: {state}, waiting...")
                time.sleep(10)
    except Exception as e:
        logging.error(f"Error checking task status: {e}")
        return None
def download_and_unzip_file(zip_url, file_name):
    """下载 ZIP 文件并解压到本地"""
    try:
        # 下载 ZIP 文件
        r = requests.get(zip_url)
        if r.status_code != 200:
            logging.error(f"Failed to download ZIP file. Status code: {r.status_code}")
            return None

        # 创建目标文件夹
        folder_path = os.path.join(download_folder, file_name[:-4])
        os.makedirs(folder_path, exist_ok=True)

        # 保存 ZIP 文件到本地
        zip_file_path = os.path.join(folder_path, file_name[:-4])
        with open(zip_file_path, "wb") as file:
            file.write(r.content)
        logging.info(f"ZIP file saved to: {zip_file_path}")

        # 解压 ZIP 文件
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            zip_ref.extractall(folder_path)
        logging.info(f"ZIP file extracted to: {folder_path}")

        # 删除临时 ZIP 文件（可选）
        os.remove(zip_file_path)
        logging.info(f"Temporary ZIP file removed: {zip_file_path}")
        
        # 复制full.md
        shutil.copyfile(os.path.join(folder_path, "full.md"), os.path.join(folder_path, f"{file_name[:-4]}.md"))

        return os.path.join(folder_path, f"{file_name[:-4]}.md")
    except Exception as e:
        logging.error(f"Error downloading or unzipping file: {e}")
        return None

def process_table_in_markdown(file_path):
    """处理 Markdown文件中的表格，替换为Typora方便阅读的格式"""
    try:
        with open(file_path, "r", encoding="utf-8") as file:
            md_text = file.read()

        md_text = md_text.replace("<html><body>", "")
        md_text = md_text.replace("</body></html>", "")

        with open(file_path, "w+", encoding="utf-8") as file:
            file.write(md_text)
        logging.info(f"Table processed and Markdown file updated: {file_path}")
        return file_path
    except Exception as e:
        logging.error(f"Error processing table in Markdown file: {e}")
        return None

# 提取翻译结果的函数
def extract_translation(res):
    # 定义正则表达式
    patterns = [
        r"(?:^|\n)\s*Translate:\s*(.+)",  # 匹配 "Translate:" 开头的结果
        r"(?:^|\n)\s*翻译:\s*(.+)",      # 匹配 "翻译:" 开头的结果
        r"<translation>(.+?)</translation>",  # 匹配 XML 格式
        r"<div[^>]*>(.+?)</div>",         # 匹配 HTML 格式
        r"```(?:json|markdown)\n(.+?)\n```",  # 匹配 Markdown 中的代码块
        r"\{\s*\"translated_text\"\s*:\s*\"(.+?)\"\s*\}",  # 匹配 JSON 格式
        r"(?:以下是翻译结果|The translation is):\s*(.+)",  # 匹配自然语言描述
        r"^(.+)$"                       # 匹配直接作为翻译结果的内容
    ]
    
    for pattern in patterns:
        match = re.search(pattern, res, re.MULTILINE | re.DOTALL)
        if match:
            extracted = match.group(1)
            
            # 如果是 JSON 格式，进一步解析
            if pattern.startswith(r"\{\s*\"translated_text\""):
                try:
                    extracted = json.loads('{' + extracted + '}')["translated_text"]
                except (json.JSONDecodeError, KeyError):
                    pass
            
            return extracted
def send_message(text, target_language="简体中文", term_mapping=None):
    """
    调用 LLM 接口完成翻译任务。
    
    参数:
        text (str): 待翻译的文本。
        target_language (str): 目标语言，默认为简体中文。
        term_mapping (dict): 专业名词对照表，格式为 {"英文术语": "中文术语"}。
    
    返回:
        str: 翻译后的文本。
    """
    # 构建专业名词对照表提示
    if term_mapping:
        term_mapping_prompt = "\n".join([f"- {key}: {value}" for key, value in term_mapping.items()])
    else:
        term_mapping_prompt = "无"

    # 构建翻译提示词
    prompt = f"""
请将待翻译文本翻译成{target_language}。翻译时需遵循以下规则：
1. 待翻译文本为markdown格式，尽量保留markdown的格式，不要删除#号和换行符。
1. 保留句子中的参考文献引用（如 [1]、(Author et al., 2023)），或句子是参考文献的引文( 如：1. Wang, S., Brose, U. & Gravel, D. Intraguild predation enhances biodiversity and functioning in complex food webs. Ecology 100, e02616 (2019).)，不要翻译、修改或删除。
2. 保留物种名、化合物名词和专业名词的原始形式（如 Escherichia coli、C6H12O6、GC-MS）。
3. 不要翻译人名、html内容、latex公式、电子邮箱
4. 使用以下专业名词对照表进行翻译：
   {term_mapping_prompt}
5. 不得遗漏或修改任何信息，尽可能还原语义。
6. 翻译结果应高质量、自然流畅，与原文的风格保持一致。
7. 直接返回翻译结果，无需任何多余的输出。

待翻译文本：
{text}
"""

    try:
        # 调用 LLM 接口
        completion = client.chat.completions.create(
            model=config.get('use_model', "qwen-plus"),  # 模型名称，可按需更换
            messages=[
                {'role': 'system', 'content': '你是一个专业的翻译专家。'},
                {'role': 'user', 'content': prompt}
            ]
        )
        
        # 提取翻译结果
        translated_text = completion.choices[0].message.content.strip()
        return extract_translation(translated_text)
    
    except Exception as e:
        logging.info(f"Error occurred while calling the LLM API: {e}")
        return None

# 读取Markdown文件
def read_markdown(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.readlines()
        
# 分类Markdown行
def classify_line(line):
    # 定义正则表达式来匹配Markdown元素
    HEADER_PATTERN = re.compile(r'^(#+)\s+(.*)')
    LIST_PATTERN = re.compile(r'^(\s*)[-*+]\s+(.*)')
    HTML_PATTERN = re.compile(r'^<[^>]+>.*<\/[^>]+>$')
    IMAGE_PATTERN = re.compile(r'!\[.*\]\(.*\)')
    if HEADER_PATTERN.match(line):
        return 'header'
    elif LIST_PATTERN.match(line):
        return 'list'
    elif HTML_PATTERN.match(line):
        return 'html'
    elif IMAGE_PATTERN.match(line):
        return 'image'
    else:
        return 'paragraph'

# 合并段落
def merge_paragraphs(lines, MAX_LENGTH = 8000):
    merged_blocks = []
    ignore_indices = []
    current_block = []
    current_length = 0
    for line in lines:
        line_type = classify_line(line)
        line_length = len(line)
        # 如果块是标题、段落、列表，可以进行下一步操作
        if line_type in ['header', 'paragraph', 'list']:
            # 如果单个块长度已经超过了最大长度，对单个块进行句子切割
            if line_length > MAX_LENGTH:
                # 对单个块进行句子切割，如果没有句子就进行长度切割
                sentences = re.split(r'(?<=[.!?])\s+', line.strip())
                if not sentences:
                    sentences = [line[i:i+MAX_LENGTH] for i in range(0, len(line), MAX_LENGTH)]
                # 不断合并句子，但不得超过最大长度，直到将块分为几段
                for sentence in sentences:
                    if current_length + len(sentence) > MAX_LENGTH:
                        merged_blocks.append('\n'.join(current_block))
                        ignore_indices.append(False)
                        current_block = []
                        current_length = 0
                    current_block.append(sentence)
                    current_length += len(sentence)
            else:
                # 如果单个块长度没有超过最大长度，不断合并块，但不得超过最大长度
                if current_length + line_length > MAX_LENGTH:
                    merged_blocks.append('\n'.join(current_block))
                    ignore_indices.append(False)
                    current_block = []
                    current_length = 0
                current_block.append(line)
                current_length += line_length
        else:
            # 否则将目前已经合并的段落放入返回结果之中，清空合并结果然后继续
            if current_block:
                merged_blocks.append('\n'.join(current_block))
                ignore_indices.append(False)
                current_block = []
                current_length = 0
            merged_blocks.append(line)
            ignore_indices.append(True)
    # 处理最后一块
    if current_block:
        merged_blocks.append('\n'.join(current_block))
        ignore_indices.append(False)
    return merged_blocks, ignore_indices

# 翻译合并后的段落
def translate_blocks(blocks, ignore_indices):
    translated_blocks = []
    indice = 0
    for block in blocks:
        if ignore_indices[indice]:
            translated_blocks.append(block + '\n')
        else:
            translated_blocks.append(send_message(block) + '\n')
        indice += 1
    return translated_blocks

def translate_markdown(file_path):
    try:
        # 读取Markdown文件
        lines = read_markdown(file_path)
        # 合并段落
        merged_blocks, ignore_indices = merge_paragraphs(lines)
        # 翻译合并后的段落
        result_lines = translate_blocks(merged_blocks, ignore_indices )
        # 保存翻译后的文件
        directory = os.path.dirname(file_path)
        file_name = os.path.basename(file_path)
        name, ext = os.path.splitext(file_name)
        new_file_name = f"译文-{name}{ext}"
        translated_file_path = os.path.join(directory, new_file_name)
        with open(translated_file_path, 'w+', encoding='utf-8') as file:
            file.writelines(result_lines)
        logging.info(f"Translated file saved to: {translated_file_path}")
        return translated_file_path
    except Exception as e:
        logging.error(f"Error translating Markdown file: {e}")
        return None

def process_pdf(file):
    try:
        logging.info("Step 1: Uploading PDF file...")
        task_id = upload_pdf(file)
        if not task_id:
            return "Failed to upload PDF file. Please check the logs for details."
        
        logging.info("Step 2: Checking task status...")
        md_url = check_task_status(task_id)
        if not md_url:
            return "Failed to retrieve Markdown file link. Please check the logs for details."
        
        logging.info("Step 3: Downloading and saving Markdown file...")
        markdown_file_path = download_and_unzip_file(md_url, os.path.basename(file))
        if not markdown_file_path:
            return "Failed to download and save Markdown file. Please check the logs for details."
        
        logging.info("Step 4: Processing table in Markdown file...")
        markdown_file_path = process_table_in_markdown(markdown_file_path)
        if not markdown_file_path:
            return "Failed to process table in Markdown file. Please check the logs for details."
        
        logging.info("Step 5: Translating Markdown file...")
        translated_file_path = translate_markdown(markdown_file_path)
        if not translated_file_path:
            return "Failed to translate Markdown file. Please check the logs for details."
        
        logging.info(f"Process completed successfully!\nTranslated file saved at: {translated_file_path}")
        return f"Process completed successfully!\nTranslated file saved at: {translated_file_path}"
    
    except Exception as e:
        os.remove(file)
        os.removedirs(os.path.dirname(file))
        logging.error(f"Unexpected error: {e}")
        return f"An unexpected error occurred: {e}. Please check the logs for details."

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from concurrent.futures import ThreadPoolExecutor

app = FastAPI()

# 定义静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")
# 结果文件目录
app.mount("/results", StaticFiles(directory="results"), name="results")

# 处理任务的线程池
executor = ThreadPoolExecutor(max_workers=4)

@app.post("/upload/")
async def upload_file(file: UploadFile = File(...)):
    try:
        # 检查文件夹是否已经存在
        folder_name = file.filename[:-4]  # 假设文件名以扩展名结尾，例如 ".pdf"
        target_folder = os.path.join(download_folder, folder_name)
        
        if os.path.exists(target_folder):
            return JSONResponse(
                content={"message": "文件已存在，请重命名文件或删除results文件夹下的同名文件夹"},
                status_code=400
            )
        
        # 创建目标文件夹
        os.makedirs(target_folder, exist_ok=True)
        
        # 保存上传的文件到本地
        file_path = os.path.join(target_folder, file.filename)
        with open(file_path, "wb") as f:
            f.write(await file.read())
        
        logging.info(f"File uploaded: {file.filename}")
        
        # 异步处理文件（假设 process_pdf 是一个处理函数）
        future = executor.submit(process_pdf, file_path)
        
        return JSONResponse(content={
            "message": "File uploaded successfully.",
            "task_status": "processing"
        })
    
    except Exception as e:
        logging.error(f"Error uploading file: {e}")
        raise HTTPException(status_code=500, detail="文件上传失败！")

@app.get("/status/{filename}")
async def check_status(filename: str):
    try:
        progress = 0
        # 检查处理后的文件是否存在
        processed_file_1 = os.path.join(download_folder, filename[:-4], f"{filename}")
        if os.path.exists(processed_file_1):
            progress = 25
        processed_file_2 = os.path.join(download_folder, filename[:-4], f"{filename[:-4]}.md")
        if os.path.exists(processed_file_2):
            progress = 50
        processed_file_3 = os.path.join(download_folder, filename[:-4], f"译文-{filename[:-4]}.md")
        if os.path.exists(processed_file_3):
            progress = 100
        if progress == 100:
            return JSONResponse(content={"status": "completed", "download_url": processed_file_3, "progress" : progress})
        else:
            return JSONResponse(content={"status": "processing", "progress" : progress})
    except Exception as e:
        logging.error(f"Error checking status: {e}")
        raise HTTPException(status_code=500, detail="Failed to check status.")


@app.get("/download/{filename}")
async def download_file(filename: str):
    try:
        # 返回处理后的文件
        processed_file = os.path.join('results', filename[:-4], f"译文-{filename[:-4]}.md")
        if os.path.exists(processed_file):
            return FileResponse(processed_file, filename=f"译文-{filename}")
        else:
            raise HTTPException(status_code=404, detail="File not found.")
    except Exception as e:
        logging.error(f"Error downloading file: {e}")
        raise HTTPException(status_code=500, detail="Failed to download file.")

